import os
import sys
from PyQt6.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QLabel, 
                             QPushButton, QFileDialog, QListWidget, QListWidgetItem, 
                             QMessageBox, QSplitter, QStatusBar, QMenu, QToolButton,
                             QFrame, QSizePolicy, QApplication, QProgressBar)
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer, QEvent, QSize, QMimeData
from PyQt6.QtGui import QIcon, QPixmap, QAction, QColor, QPalette, QDragEnterEvent, QDropEvent
import cv2
from datetime import datetime
import shutil

# 导入自定义参数面板
from .waifu2x_caffe_params import Waifu2xCaffeParams

class ProgressBar(QWidget):
    """自定义进度条，显示当前处理进度"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        
        # 创建主布局
        self.main_layout = QVBoxLayout()
        self.main_layout.setContentsMargins(10, 5, 10, 5)
        
        # 创建进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setTextVisible(True)
        self.progress_bar.setFormat("%p%")
        
        # 添加到布局
        self.main_layout.addWidget(self.progress_bar)
        
        # 设置布局
        self.setLayout(self.main_layout)
    
    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)

class StatusBarWithProgress(QStatusBar):
    """带有状态信息的状态栏"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        
        # 创建主布局
        self.main_layout = QHBoxLayout()
        self.main_layout.setContentsMargins(0, 0, 0, 0)
        
        # 创建状态标签
        self.status_label = QLabel("准备就绪")
        self.status_label.setMinimumWidth(200)
        
        # 创建文件处理标签
        self.file_label = QLabel("")
        
        # 添加到布局
        self.main_layout.addWidget(self.status_label)
        self.main_layout.addWidget(self.file_label)
        
        # 创建一个容器来放置我们的布局
        container = QWidget()
        container.setLayout(self.main_layout)
        
        # 添加容器到状态栏
        self.addPermanentWidget(container, 1)
    
    def update_status(self, text=""):
        """更新状态文本"""
        if text:
            self.status_label.setText(text)
    
    def update_file_progress(self, filename):
        """更新当前处理的文件名"""
        if filename:
            self.file_label.setText(f"处理: {filename}")
        else:
            self.file_label.setText("")

class FileListWidget(QListWidget):
    """带有拖放功能的文件列表控件"""
    
    files_dropped = pyqtSignal(list)  # 文件拖放信号
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setAcceptDrops(True)  # 启用接受拖放
        self.setSelectionMode(QListWidget.SelectionMode.ExtendedSelection)  # 允许多选
    
    def dragEnterEvent(self, event: QDragEnterEvent):
        """拖入事件"""
        # 检查是否包含文件数据
        if event.mimeData().hasUrls():
            # 接受拖放事件
            event.accept()
    
    def dragMoveEvent(self, event):
        """拖动事件"""
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()
    
    def dropEvent(self, event: QDropEvent):
        """放置事件"""
        file_paths = []
        
        if event.mimeData().hasUrls():
            event.accept()
            
            for url in event.mimeData().urls():
                # 获取本地文件路径
                file_path = url.toLocalFile()
                # 检查是否是图像文件
                if self.is_image_file(file_path):
                    file_paths.append(file_path)
                    # 添加到列表控件
                    self.add_file(file_path)
            
            # 如果有有效文件，发出信号
            if file_paths:
                self.files_dropped.emit(file_paths)
        else:
            event.ignore()
    
    def is_image_file(self, file_path):
        """检查文件是否是图像"""
        image_extensions = [".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff"]
        _, ext = os.path.splitext(file_path.lower())
        return ext in image_extensions
    
    def add_file(self, file_path):
        """添加文件到列表"""
        # 检查是否已存在
        for i in range(self.count()):
            if self.item(i).data(Qt.ItemDataRole.UserRole) == file_path:
                return  # 已存在，不添加
        
        # 创建列表项
        item = QListWidgetItem(os.path.basename(file_path))
        item.setData(Qt.ItemDataRole.UserRole, file_path)  # 存储完整路径
        self.addItem(item)
    
    def get_selected_files(self):
        """获取选中的文件路径列表"""
        files = []
        for item in self.selectedItems():
            files.append(item.data(Qt.ItemDataRole.UserRole))
        return files
    
    def get_all_files(self):
        """获取所有文件路径列表"""
        files = []
        for i in range(self.count()):
            files.append(self.item(i).data(Qt.ItemDataRole.UserRole))
        return files
    
    def remove_selected_files(self):
        """移除选中的文件"""
        for item in self.selectedItems():
            row = self.row(item)
            self.takeItem(row)

class MainWindow(QMainWindow):
    """Waifu2x-Caffe图像超分辨率工具主窗口"""
    
    def __init__(self, image_processor):
        super().__init__()
        self.image_processor = image_processor
        self.default_save_dir = ""
        
        # 初始化UI
        self.init_ui()
        
        # 应用暗色主题
        self.apply_dark_theme()
        
        # 连接信号
        self.connect_signals()
    
    def init_ui(self):
        """初始化UI界面"""
        # 创建中央部件和布局
        central_widget = QWidget()
        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(10, 10, 10, 10)
        main_layout.setSpacing(10)  # 增加布局间距
        
        # 创建状态栏
        self.status_bar = StatusBarWithProgress()
        self.setStatusBar(self.status_bar)
        
        # 创建主分割器（左右布局）
        self.main_splitter = QSplitter(Qt.Orientation.Horizontal)
        self.main_splitter.setHandleWidth(1)  # 设置分割线宽度
        
        # 左侧：文件列表区域
        self.create_file_list_area()
        
        # 右侧：参数设置区域
        self.create_params_area()
        
        # 设置主分割器初始比例
        self.main_splitter.setSizes([300, 500])
        
        # 添加主分割器到布局
        main_layout.addWidget(self.main_splitter, 1)  # 让分割器占据大部分空间
        
        # 创建进度条
        self.progress_area = ProgressBar()
        main_layout.addWidget(self.progress_area, 0)  # 进度条不拉伸
        
        # 设置中央部件
        self.setCentralWidget(central_widget)
        
        # 创建菜单栏
        self.create_menu_bar()
        
        # 设置窗口标题和大小
        self.setWindowTitle('Waifu2x-Caffe 图像超分辨率工具')
        self.resize(1000, 650)
        
        # 启用整个窗口的拖放
        self.setAcceptDrops(True)
    
    def create_menu_bar(self):
        """创建菜单栏"""
        # 菜单栏
        menubar = self.menuBar()
        
        # 文件菜单
        file_menu = menubar.addMenu('文件')
        
        # 添加图像
        add_image_action = QAction('添加图像', self)
        add_image_action.triggered.connect(self.on_add_files)
        file_menu.addAction(add_image_action)
        
        # 清空缓存数据
        clear_cache_action = QAction('清空缓存数据', self)
        clear_cache_action.triggered.connect(self.on_clear_cache)
        file_menu.addAction(clear_cache_action)
        
        file_menu.addSeparator()
        
        # 退出
        exit_action = QAction('退出', self)
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)
        
        # 处理菜单
        process_menu = menubar.addMenu('处理')
        
        # 处理图像
        process_action = QAction('处理图像', self)
        process_action.triggered.connect(self.on_process_images)
        process_menu.addAction(process_action)
    
    def create_file_list_area(self):
        """创建文件列表区域"""
        # 文件列表容器
        file_list_container = QWidget()
        file_layout = QVBoxLayout(file_list_container)
        file_layout.setContentsMargins(10, 10, 10, 10)
        file_layout.setSpacing(8)  # 设置垂直间距
        
        # 文件列表标题
        title_layout = QHBoxLayout()
        title_label = QLabel("图像文件列表")
        title_label.setStyleSheet("font-weight: bold; font-size: 14px;")
        title_layout.addWidget(title_label)
        
        # 添加/移除按钮
        button_layout = QHBoxLayout()
        button_layout.setSpacing(5)  # 设置按钮间的水平间距
        
        self.add_files_btn = QPushButton("添加文件")
        self.add_files_btn.setMinimumWidth(90)  # 设置最小宽度
        self.add_files_btn.clicked.connect(self.on_add_files)
        button_layout.addWidget(self.add_files_btn)
        
        self.remove_files_btn = QPushButton("移除选中")
        self.remove_files_btn.setMinimumWidth(90)  # 设置最小宽度
        self.remove_files_btn.clicked.connect(self.on_remove_selected)
        button_layout.addWidget(self.remove_files_btn)
        
        self.clear_files_btn = QPushButton("清空列表")
        self.clear_files_btn.setMinimumWidth(90)  # 设置最小宽度
        self.clear_files_btn.clicked.connect(self.on_clear_files)
        button_layout.addWidget(self.clear_files_btn)
        
        # 文件列表控件
        self.file_list = FileListWidget()
        self.file_list.setMinimumWidth(200)
        
        # 拖放提示标签
        self.drag_label = QLabel("将图像文件拖放到此处")
        self.drag_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.drag_label.setStyleSheet("color: #777777; font-style: italic;")
        
        # 添加到布局
        file_layout.addLayout(title_layout)
        file_layout.addLayout(button_layout)
        file_layout.addWidget(self.file_list)
        file_layout.addWidget(self.drag_label)
        
        # 添加到主分割器
        self.main_splitter.addWidget(file_list_container)
    
    def create_params_area(self):
        """创建参数设置区域"""
        # 参数设置容器
        params_container = QWidget()
        params_layout = QVBoxLayout(params_container)
        params_layout.setContentsMargins(10, 10, 10, 10)
        params_layout.setSpacing(10)  # 设置垂直间距
        
        # 参数标题
        params_title = QLabel("处理参数设置")
        params_title.setStyleSheet("font-weight: bold; font-size: 14px;")
        params_layout.addWidget(params_title)
        
        # 参数面板
        self.params_panel = Waifu2xCaffeParams()
        params_layout.addWidget(self.params_panel)
        
        # 处理按钮区域
        buttons_layout = QHBoxLayout()
        buttons_layout.setContentsMargins(0, 10, 0, 0)  # 增加与参数面板的距离
        
        self.process_btn = QPushButton("处理图像")
        self.process_btn.setMinimumWidth(120)  # 设置最小宽度
        self.process_btn.setMinimumHeight(30)  # 设置最小高度
        self.process_btn.clicked.connect(self.on_process_images)
        buttons_layout.addWidget(self.process_btn)
        
        params_layout.addLayout(buttons_layout)
        params_layout.addStretch(1)  # 添加弹性空间
        
        # 添加到主分割器
        self.main_splitter.addWidget(params_container)
    
    def connect_signals(self):
        """连接信号槽"""
        # 文件列表拖放信号
        self.file_list.files_dropped.connect(self.on_files_dropped)
        
        # 图像处理信号
        self.image_processor.progress_updated.connect(self.on_progress_updated)
        self.image_processor.file_progress_updated.connect(self.on_file_progress_updated)
        self.image_processor.process_finished.connect(self.on_process_finished)
        self.image_processor.batch_finished.connect(self.on_batch_finished)
        self.image_processor.process_error.connect(self.on_process_error)
    
    def on_progress_updated(self, value, text=""):
        """进度更新响应"""
        self.progress_area.update_progress(value)
        if text:
            self.status_bar.update_status(text)
    
    def on_file_progress_updated(self, filename, progress):
        """文件进度更新响应"""
        self.status_bar.update_file_progress(filename)
        self.progress_area.update_progress(progress)
    
    # 整个窗口的拖放事件
    def dragEnterEvent(self, event):
        """窗口拖入事件"""
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()
    
    def dragMoveEvent(self, event):
        """窗口拖动事件"""
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()
    
    def dropEvent(self, event):
        """窗口放置事件"""
        if event.mimeData().hasUrls():
            event.accept()
            file_paths = []
            
            for url in event.mimeData().urls():
                # 获取本地文件路径
                file_path = url.toLocalFile()
                # 检查是否是图像文件
                if self.file_list.is_image_file(file_path):
                    file_paths.append(file_path)
                    # 添加到列表控件
                    self.file_list.add_file(file_path)
            
            # 如果有有效文件，发出信号
            if file_paths:
                self.on_files_dropped(file_paths)
        else:
            event.ignore()
    
    def apply_dark_theme(self):
        """应用暗色主题"""
        self.setStyleSheet("""
            QMainWindow, QWidget {
                background-color: #2D2D30;
                color: #CCCCCC;
            }
            QMenuBar {
                background-color: #1E1E1E;
                color: #CCCCCC;
            }
            QMenuBar::item:selected {
                background-color: #3E3E40;
            }
            QMenu {
                background-color: #1E1E1E;
                color: #CCCCCC;
                border: 1px solid #555555;
            }
            QMenu::item:selected {
                background-color: #3E3E40;
            }
            QPushButton {
                background-color: #3E3E40;
                color: #CCCCCC;
                border: 1px solid #555555;
                padding: 5px 15px;
                border-radius: 2px;
            }
            QPushButton:hover {
                background-color: #505050;
            }
            QPushButton:pressed {
                background-color: #0078D7;
            }
            QListWidget {
                background-color: #252526;
                color: #CCCCCC;
                border: 1px solid #3F3F46;
                border-radius: 2px;
            }
            QListWidget::item:selected {
                background-color: #0078D7;
            }
            QSplitter::handle {
                background-color: #3F3F46;
            }
            QLabel {
                color: #CCCCCC;
            }
            QProgressBar {
                border: 1px solid #3F3F46;
                background-color: #252526;
                border-radius: 2px;
                text-align: center;
            }
            QProgressBar::chunk {
                background-color: #0078D7;
                border-radius: 2px;
            }
            QComboBox, QLineEdit {
                background-color: #252526;
                color: #CCCCCC;
                border: 1px solid #3F3F46;
                border-radius: 2px;
                padding: 3px;
            }
            QGroupBox {
                border: 1px solid #3F3F46;
                border-radius: 4px;
                margin-top: 8px;
                padding-top: 10px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                subcontrol-position: top center;
                padding: 0 3px;
            }
            QCheckBox, QRadioButton {
                color: #CCCCCC;
                spacing: 5px;
            }
            QStatusBar {
                background-color: #1E1E1E;
                color: #CCCCCC;
                border-top: 1px solid #3F3F46;
            }
        """)
    
    def on_add_files(self):
        """添加文件按钮点击响应"""
        files, _ = QFileDialog.getOpenFileNames(
            self,
            "选择图像文件",
            "",
            "图像文件 (*.png *.jpg *.jpeg *.bmp *.webp *.tiff)"
        )
        
        if files:
            for file_path in files:
                self.file_list.add_file(file_path)
            
            # 如果添加了文件，隐藏拖放提示
            if self.file_list.count() > 0:
                self.drag_label.setVisible(False)
    
    def on_files_dropped(self, file_paths):
        """文件拖放响应"""
        # 文件已在FileListWidget中添加到列表，此处只需显示提示
        self.status_bar.update_status(f"已添加 {len(file_paths)} 个文件")
        
        # 如果列表中有文件，隐藏拖放提示
        if self.file_list.count() > 0:
            self.drag_label.setVisible(False)
    
    def on_remove_selected(self):
        """移除选中文件按钮响应"""
        self.file_list.remove_selected_files()
        
        # 如果列表为空，显示拖放提示
        if self.file_list.count() == 0:
            self.drag_label.setVisible(True)
    
    def on_clear_files(self):
        """清空文件列表按钮响应"""
        self.file_list.clear()
        
        # 显示拖放提示
        self.drag_label.setVisible(True)
    
    def on_clear_cache(self):
        """清空缓存数据"""
        # 确认对话框
        reply = QMessageBox.question(
            self,
            "清空缓存确认",
            "确定要清空Waifu2x-Caffe缓存数据吗？\n这可能会导致下次处理时间稍长。",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No
        )
        
        if reply == QMessageBox.StandardButton.Yes:
            # 清空缓存路径
            cache_dir = os.path.join(
                os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                'models',
                'waifu2x-caffe',
                'waifu2x-caffe-cudnn_data'
            )
            
            try:
                if os.path.exists(cache_dir):
                    # 删除目录及其所有内容
                    shutil.rmtree(cache_dir)
                    self.status_bar.update_status("缓存数据已清空")
                    QMessageBox.information(self, "清空成功", "Waifu2x-Caffe缓存数据已成功清空")
                else:
                    QMessageBox.information(self, "清空成功", "缓存目录不存在或已经是空的")
            except Exception as e:
                QMessageBox.critical(self, "清空失败", f"清空缓存数据时出错: {str(e)}")
    
    def on_process_images(self):
        """处理图像按钮响应"""
        # 获取文件列表数量
        count = self.file_list.count()
        
        # 如果文件列表为空，提示用户添加文件
        if count == 0:
            QMessageBox.warning(self, "提示", "请先添加要处理的图像文件")
            return
        
        # 根据列表中文件数量决定处理方式
        if count == 1:
            # 单个文件处理
            file_path = self.file_list.item(0).data(Qt.ItemDataRole.UserRole)
            self.process_images([file_path])
        else:
            # 获取选中的文件
            selected_files = self.file_list.get_selected_files()
            
            # 如果有选中文件，处理选中的；否则处理所有的
            if selected_files:
                self.process_images(selected_files)
            else:
                # 获取所有文件
                all_files = self.file_list.get_all_files()
                self.process_images(all_files)
    
    def process_images(self, file_paths):
        """处理图像"""
        # 获取参数
        params = self.params_panel.get_params()
        scale_factor = params.get('scale', 2)
        
        # 如果多于一个文件，确认批量处理
        if len(file_paths) > 1:
            # 批量处理所有图像
            self.status_bar.update_status(f"开始批量处理 {len(file_paths)} 个文件...")
            self.image_processor.batch_process(file_paths, None, params)
        else:
            # 处理单个图像
            image_path = file_paths[0]
            self.status_bar.update_status(f"正在处理: {os.path.basename(image_path)}...")
            self.image_processor.process_image(image_path, params)
    
    def on_process_finished(self, result_data):
        """单个图像处理完成响应"""
        if not isinstance(result_data, dict) or 'image' not in result_data:
            QMessageBox.warning(self, "处理失败", "图像处理失败")
            return
        
        result_image = result_data['image']
        original_path = result_data['original_path']
        scale_factor = result_data['scale_factor']
        
        # 自动保存至原图片目录
        dir_name = os.path.dirname(original_path)
        base_name, ext = os.path.splitext(os.path.basename(original_path))
        output_path = os.path.join(dir_name, f"{base_name}_{scale_factor}x{ext}")
        
        # 保存图像
        cv2.imwrite(output_path, result_image)
        self.status_bar.update_status(f"处理结果已保存到: {output_path}")
        self.progress_area.update_progress(100)
        
        # 显示保存成功消息
        QMessageBox.information(
            self, 
            "处理完成", 
            f"处理完成，结果已保存至:\n{output_path}"
        )
    
    def on_batch_finished(self, result_files):
        """批量处理完成响应"""
        if not result_files:
            QMessageBox.warning(self, "处理结果", "批量处理失败或无任何结果")
            return
        
        QMessageBox.information(
            self, 
            "批量处理完成", 
            f"成功处理 {len(result_files)} 个文件，已保存到原始图像所在目录"
        )
        self.status_bar.update_status(f"批量处理完成，处理了 {len(result_files)} 个文件")
        self.progress_area.update_progress(100)
    
    def on_process_error(self, error_msg):
        """处理错误响应"""
        QMessageBox.critical(self, "处理错误", error_msg)
        self.status_bar.update_status(f"处理失败: {error_msg}")
        self.progress_area.update_progress(0)
